package middleware

import (
	"time"

	"gitee.com/krio/ginc/example/core/exerr"
	redisLib "gitee.com/krio/ginc/example/core/lib/redis"
	"gitee.com/krio/ginc/example/router/output"
	"github.com/gin-gonic/gin"
	"github.com/gomodule/redigo/redis"
	ginLimitByKey "github.com/yangxikun/gin-limit-by-key"
	"golang.org/x/time/rate"
)

const LimiterMiddlewareKey = "Limiter" // 中间件标识

type LimitType string // 限流器类型

const (
	ReteLimit LimitType = "ReteLimit"
	RLTBLimit LimitType = "RLTBLimit"
)

type LimitMode string // 限流模式

const (
	LimitIp        LimitMode = "LIMIT_IP"          // 仅IP
	LimitIpAndPath LimitMode = "LIMIT_IP_AND_PATH" // IP+路由地址
	LimitUser      LimitMode = "LIMIT_USER"        // 仅用户
)

type LimitOption struct {
	Type         LimitType // 限流器类型
	Mode         LimitMode // 限流模式
	Cap          int64     // 容量数
	IntervalTime int64     // 令牌间隔时间(毫秒)
	ExpireTime   int64     // 令牌过期时间(毫秒)
}

// Limiter 接口限流器
func Limiter(opt *LimitOption) gin.HandlerFunc {
	if opt.Mode == "" {
		opt.Mode = LimitIpAndPath
	}
	if opt.Cap <= 0 {
		opt.Cap = 100
	}
	if opt.IntervalTime <= 0 {
		opt.IntervalTime = 300
	}
	if opt.ExpireTime <= 0 {
		opt.ExpireTime = 5000
	}
	switch opt.Type {
	case ReteLimit:
		return reteLimit(opt)
	case RLTBLimit:
		return rltbLimit(opt)
	default:
		return reteLimit(opt)
	}
}

func HandleLimitMode(ctx *gin.Context, mode LimitMode) string {
	switch mode {
	case LimitIp:
		return ctx.ClientIP()
	case LimitIpAndPath:
		return ctx.ClientIP() + ctx.FullPath()
	default:
		return ctx.ClientIP()
	}
}

// reteLimit 限流器 (基于rate类库 + 令牌桶算法)
func reteLimit(opt *LimitOption) gin.HandlerFunc {
	return ginLimitByKey.NewRateLimiter(func(ctx *gin.Context) string {
		return HandleLimitMode(ctx, opt.Mode)
	}, func(ctx *gin.Context) (*rate.Limiter, time.Duration) {
		return rate.NewLimiter(
				rate.Every(time.Duration(opt.IntervalTime)),
				int(opt.Cap)),
			time.Duration(opt.ExpireTime)
	}, func(ctx *gin.Context) {
		output.Response(ctx, "", exerr.FrequencyTooMany)
		ctx.Abort()
		return
	})
}

// rltbLimit 限流器 (基于redis + lua + 令牌桶算法)
func rltbLimit(opt *LimitOption) gin.HandlerFunc {
	return func(ctx *gin.Context) {
		conn := redisLib.GetRedisPool().Get()
		lua := redis.NewScript(1, LimitTokenBucket)
		arg := []interface{}{
			HandleLimitMode(ctx, opt.Mode), // key
			opt.Cap,                        // ARGV[1] bucket_cap 令牌桶容量
			opt.IntervalTime,               // ARGV[2] interval 放令牌间隔（毫秒）
			time.Now().UnixNano() / 1e6,    // ARGV[3] curr_timestamp 获取当前 (毫秒)
			opt.ExpireTime,                 // ARGV[4] expire_time 过期时间 (毫秒)
		}
		args := redis.Args{}.AddFlat(arg)
		in, err := redis.Int(lua.Do(conn, args...))
		if err != nil {
			output.Response(ctx, "", exerr.Redis.WithSubmsg(err.Error()))
			ctx.Abort()
			return
		}
		if in != 1 {
			output.Response(ctx, "", exerr.FrequencyTooMany)
			ctx.Abort()
			return
		}
		ctx.Next()
	}
}

// Lua内容
const (
	// LimitTokenBucket 令牌桶算法
	LimitTokenBucket = `
	-- 返回码：
	local SUCCESS = 1             -- 操作成功
	local FAILED = 0              -- 操作失败

	local key = KEYS[1]                       -- key
	local bucket_cap = tonumber(ARGV[1])      -- 令牌桶容量
	local interval = tonumber(ARGV[2])        -- 放令牌间隔(毫秒)
	local cur_token_count = bucket_cap      -- 目前令牌桶中的令牌数量（默认桶容量）
	local cur_timestamp = tonumber(ARGV[3])   -- 当前时间戳（毫秒）
	local expire_time = tonumber(ARGV[4])     -- 过期时间（毫秒）
	
	-- （验证请求）桶容量、令牌间隔时间，必须大于 0
    if bucket_cap <= 0 or interval <= 0 then
		return FAILED
    end
	-- 过期时间过小
	if expire_time <= interval then
		return FAILED
    end
	
	-- 桶是否存在
	local exi = redis.call("EXISTS", key)
	if exi == 0 then  -- 不存在
        redis.call("HMSET", key, 
		"last_timestamp", cur_timestamp, 
		"cur_token_count", cur_token_count - 1) -- 首次请求，默认为满桶，消耗一个 token
        -- 设置 redis 的过期时间
		redis.call("EXPIRE", key, expire_time)
		return SUCCESS
	end

	local arr = redis.pcall("HMGET", KEYS[1], "last_timestamp", "cur_token_count")
    if arr == nil then -- 空内容
        return FAILED
    end

	local last_timestamp = tonumber(arr[1])
	cur_token_count =  tonumber(arr[2]) -- 令牌桶的 token
	-- 计算上一次放令牌到现在的时间间隔中，一共应该放入多少令牌
	local reserve_count = math.max(0, math.floor((cur_timestamp - last_timestamp) / interval))
	local new_count = math.min(bucket_cap, cur_token_count + reserve_count)

	-- 无令牌可用
	if new_count <= 0 then
		-- 更新数据
		redis.call("HMSET", key, "last_timestamp", last_timestamp, "cur_token_count", new_count)
		-- 设置 redis 的过期时间
		redis.call("EXPIRE", key, expire_time)

		return FAILED
	end

	-- 更新当前桶中的令牌数量 
	-- 如果这次有放入令牌，则更新时间
	redis.pcall("HSET", key, "last_timestamp", cur_timestamp, "cur_token_count", new_count -1)
	redis.call("EXPIRE", key, expire_time) -- 设置 redis 的过期时间

	return SUCCESS`
)
